#include <wdm.h>

#include "mydriver.h"

NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING pRegistryPath)
{
	NTSTATUS status = STATUS_SUCCESS;
	PDEVICE_OBJECT p_device_object = NULL;
	PDEVICE_EXTENSION pdx = NULL;
	UNICODE_STRING driver_name, dos_driver_name;
	
	UNREFERENCED_PARAMETER(pRegistryPath);
	
	DbgPrint("DriverEntry called\r\n");
	
	RtlInitUnicodeString(&driver_name, DRVNAME);
	RtlInitUnicodeString(&dos_driver_name, DOSDRVNAME);
	
	status = IoCreateDevice(pDriverObject
							, sizeof(DEVICE_EXTENSION)
							, &driver_name
							, FILE_DEVICE_UNKNOWN
							, FILE_DEVICE_SECURE_OPEN
							, FALSE
							, &p_device_object);
	pdx = (PDEVICE_EXTENSION) p_device_object->DeviceExtension;
	pdx->memory_buffer = NULL;
	pdx->buffer_length = 0;

	if (!NT_SUCCESS(status)) return status;

	pDriverObject->MajorFunction[IRP_MJ_CREATE] = CreateMyDriver;
	pDriverObject->MajorFunction[IRP_MJ_CLOSE] = CloseMyDriver;
	pDriverObject->MajorFunction[IRP_MJ_WRITE] = WriteMyDriver;
	pDriverObject->MajorFunction[IRP_MJ_READ] = ReadMyDriver;
	pDriverObject->DriverUnload = UnloadMyDriver;
	
	p_device_object->Flags |= DO_DIRECT_IO;
	p_device_object->Flags &= (~DO_DEVICE_INITIALIZING);
	
	status = IoCreateSymbolicLink(&dos_driver_name, &driver_name);
	
	if (!NT_SUCCESS(status)) IoDeleteDevice(pDriverObject->DeviceObject);
	
	return status;
}

NTSTATUS CreateMyDriver(PDEVICE_OBJECT pDeviceObject, PIRP pIrp)
{
	UNREFERENCED_PARAMETER(pDeviceObject);
	UNREFERENCED_PARAMETER(pIrp);
	PAGED_CODE();

	DbgPrint("CreateMyDriver called\r\n");
	
	return STATUS_SUCCESS;
}

NTSTATUS CloseMyDriver(PDEVICE_OBJECT pDeviceObject, PIRP pIrp)
{
	UNREFERENCED_PARAMETER(pDeviceObject);
	UNREFERENCED_PARAMETER(pIrp);
	PAGED_CODE();

	DbgPrint("CloseMyDriver called\r\n");
	
	return STATUS_SUCCESS;
}

NTSTATUS WriteMyDriver(PDEVICE_OBJECT pDeviceObject, PIRP pIrp)
{
	PIO_STACK_LOCATION pIoStackIrp = NULL;
	PUCHAR pWriteDataBuffer;
	
	PAGED_CODE();
	
	DbgPrint("WriteMyDriver called\r\n");
	
	pIoStackIrp = IoGetCurrentIrpStackLocation(pIrp);
	if (pIoStackIrp)
	{
		pWriteDataBuffer = MmGetSystemAddressForMdlSafe(pIrp->MdlAddress, NormalPagePriority);
		if (pWriteDataBuffer)
		{
			PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) pDeviceObject->DeviceExtension;
			if (pdx->buffer_length != 0) ExFreePool(pdx->memory_buffer);
			pdx->buffer_length = pIoStackIrp->Parameters.Write.Length;
			pdx->memory_buffer = ExAllocatePoolWithTag(NonPagedPool, pdx->buffer_length, 'clpb'); //WTF?! character literal ?
			if (pdx->memory_buffer != NULL)
				RtlCopyMemory(pdx->memory_buffer, pWriteDataBuffer, pdx->buffer_length);
			else
				pdx->buffer_length = 0;
		}
	}
	
	return STATUS_SUCCESS;
}

NTSTATUS ReadMyDriver(PDEVICE_OBJECT pDeviceObject, PIRP pIrp)
{
	PIO_STACK_LOCATION pIoStackIrp = NULL;
	PUCHAR pReadDataBuffer;
	
	PAGED_CODE();
	
	DbgPrint("ReadMyDriver called\r\n");
	
	pIoStackIrp = IoGetCurrentIrpStackLocation(pIrp);
	if (pIoStackIrp)
	{
		PDEVICE_EXTENSION pdx = (PDEVICE_EXTENSION) pDeviceObject->DeviceExtension;
		pReadDataBuffer = MmGetSystemAddressForMdlSafe(pIrp->MdlAddress, NormalPagePriority);
		if (pReadDataBuffer && pIoStackIrp->Parameters.Read.Length >= pdx->buffer_length)
		{
			RtlCopyMemory(pReadDataBuffer, pdx->memory_buffer, pdx->buffer_length);
			if (pdx->buffer_length != 0)
			{
				ExFreePool(pdx->memory_buffer);
				pdx->memory_buffer = NULL;
				pdx->buffer_length = 0;
			}
			return STATUS_SUCCESS;
		}
	}
	
	return STATUS_BUFFER_TOO_SMALL;
}

VOID UnloadMyDriver(PDRIVER_OBJECT pDriverObject)
{
	UNICODE_STRING dos_driver_name;

	PAGED_CODE();

	DbgPrint("UnloadMyDriver called\r\n");
	RtlInitUnicodeString(&dos_driver_name, DOSDRVNAME);
	IoDeleteSymbolicLink(&dos_driver_name);
	
	IoDeleteDevice(pDriverObject->DeviceObject);
}